from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage
import dashscope
from http import HTTPStatus
from urllib.parse import urlparse, unquote
from pathlib import PurePosixPath
from dashscope import ImageSynthesis


class ChatGPT(object):
    def __init__(self, model, api_key=None, base_url=None):
        # 选择调用的OpenAI的模型
        self.model = ChatOpenAI(
            model=model,
            api_key=api_key,
            base_url=base_url,
        )

    def send(self, prompt, history=None):
        if history is None:
            history = []

        response = self.model.invoke([
            *history,
            HumanMessage(content=prompt),
        ])
        return response


class Tongyi(object):
    def __init__(self, api_key):
        self.api_key = api_key

    def text_send(self, prompt, history=None):
        """文生文"""
        if history is None:
            history = []

        messages = [
            *history,
            {'role': 'user', 'content': prompt},
        ]
        print(messages)
        response = dashscope.Generation.call(
            model=dashscope.Generation.Models.qwen_turbo,
            messages=messages,
            result_format='message',  # 将返回结果格式设置为 message
            api_key=self.api_key,
        )
        print(response)
        if response.status_code == HTTPStatus.OK:
            return response.output.choices[0].message['content']
        else:
            return None

    def img_send(self, prompt, num=4, size='720*1280'):
        response = ImageSynthesis.call(
            model=ImageSynthesis.Models.wanx_v1,  # 模型名称
            prompt=prompt,  # 提示语
            n=num,  # 图片数量
            size=size,  # 图片尺寸
            api_key=self.api_key,
        )
        if response.status_code == HTTPStatus.OK:
            data = {}
            for result in response.output.results:
                data[PurePosixPath(unquote(urlparse(result.url).path)).parts[-1]] = result.url
            return data
        else:
            return {}


if __name__ == '__main__':
    # robot = ChatGPT(
    #     model="gpt-3.5-turbo",
    #     api_key="sk-WmM426a0b2ce516cdfab7a58d6562fe2d78cbcd5288IYPXT",
    #     base_url="https://api.gptsapi.net/v1"
    # )
    #
    # response = robot.send("你是谁？")
    # print(response.content)

    ty = Tongyi(api_key='sk-c98c8f9c9e474aea9c4d56bd52fcc001')
    # # 文生文
    # response = ty.text_send("请介绍下你的功能和接口")
    # print(response)

    # # 文生图
    # response = ty.img_send('画一张水墨风格的美女图')
    # for file_name, url in response.items():
    #     with open('./%s' % file_name, 'wb+') as f:
    #         f.write(requests.get(url).content)
